# update_requirements.py
import subprocess
import sys
import re
import packaging.version # Import packaging library for version comparison

def get_latest_version(package_name):
    """Queries pip index for the latest version of a package."""
    try:
        # Run pip index versions command
        result = subprocess.run(
            [sys.executable, '-m', 'pip', 'index', 'versions', package_name],
            capture_output=True,
            text=True,
            check=True,
            encoding='utf-8' # Specify encoding
        )

        # --- Modified regexes to handle various formats ---
        # Search for the "LATEST:" line
        latest_match = re.search(r"^\s*LATEST:\s*([0-9a-zA-Z\.\-]+)", result.stdout, re.MULTILINE)
        if latest_match:
            return latest_match.group(1)

        # Search for version in parentheses on the first line (common format)
        first_line_match = re.search(r"^[a-zA-Z0-9_\-]+ \(([0-9a-zA-Z\.\-]+)\)", result.stdout)
        if first_line_match:
            return first_line_match.group(1)

        # Fallback to previous regexes if above formats are not found
        latest_match_fallback1 = re.search(r"Latest version:\s*([0-9a-zA-Z\.\-]+)", result.stdout)
        if latest_match_fallback1:
            return latest_match_fallback1.group(1)

        versions_match_fallback2 = re.findall(r"([0-9a-zA-Z\.\-]+)\s+\(latest\)", result.stdout)
        if versions_match_fallback2:
             return versions_match_fallback2[0] # Assuming the first one marked (latest) is it
        # --- End of modifications ---


        print(f"警告：无法从输出中找到 {package_name} 的最新版本信息：\n{result.stdout}", file=sys.stderr) # Adjusted message
        return None

    except subprocess.CalledProcessError as e:
        print(f"查询 {package_name} 最新版本信息时出错： {e}", file=sys.stderr) # Translated message
        print(f"标准错误输出： {e.stderr}", file=sys.stderr) # Translated message
        print(f"标准输出： {e.stdout}", file=sys.stderr) # Translated message
        return None
    except Exception as e:
        print(f"查询 {package_name} 时发生意外错误： {e}", file=sys.stderr) # Translated message
        return None


def update_requirements(requirements_file="requirements.txt"):
    """Reads requirements.txt, updates versions, and writes a new file."""
    updated_lines = []
    try:
        with open(requirements_file, 'r', encoding='utf-8') as f:
            lines = f.readlines()

        print(f"正在读取 {requirements_file}...") # Translated message

        for line in lines:
            original_line = line.strip()
            if not original_line or original_line.startswith('#'):
                updated_lines.append(original_line) # Keep comments and blank lines
                continue

            # Extract package name and potential current version
            match = re.match(r"^\s*([a-zA-Z0-9_\-]+)(?:==)?([0-9a-zA-Z\.\-]+)?(?:\s*#.*)?$", original_line)
            if not match:
                 print(f"警告：无法从行 '{original_line}' 解析包名称或版本。跳过该行。", file=sys.stderr) # Adjusted message
                 updated_lines.append(original_line)
                 continue

            package_name = match.group(1)
            current_version_str = match.group(2) # Group 2 is the version string if '==' was used

            print(f"正在检查 {package_name} 的最新版本...") # Translated message
            latest_version = get_latest_version(package_name)

            if latest_version:
                if current_version_str:
                    try:
                        # Compare versions using packaging library
                        current_ver = packaging.version.parse(current_version_str)
                        latest_ver = packaging.version.parse(latest_version)

                        if latest_ver > current_ver:
                            # Update the version
                            comment_match = re.search(r"(\s*#.*)$", original_line)
                            comment = comment_match.group(1) if comment_match else ""
                            new_line = f"{package_name}=={latest_version}{comment}"
                            print(f"  找到最新版本: {latest_version}. 正在将 '{original_line}' 更新为 '{new_line}'") # Translated message
                            updated_lines.append(new_line)
                        else:
                            # Already latest or newer (e.g., local build) - keep original
                            print(f"  {package_name} ({current_version_str}) 已是最新版本。", file=sys.stderr) # New message
                            updated_lines.append(original_line)

                    except packaging.version.InvalidVersion as e:
                        print(f"警告：解析 {package_name} 的版本号 '{current_version_str}' 时出错：{e}. 保留原始行：'{original_line}'", file=sys.stderr) # New message
                        updated_lines.append(original_line)
                    except Exception as e:
                         print(f"警告：比较 {package_name} 的版本时发生意外错误：{e}. 保留原始行：'{original_line}'", file=sys.stderr) # New message
                         updated_lines.append(original_line)

                else: # No version specified in requirements.txt line
                     # Just add the latest version with ==
                    comment_match = re.search(r"(\s*#.*)$", original_line)
                    comment = comment_match.group(1) if comment_match else ""
                    new_line = f"{package_name}=={latest_version}{comment}"
                    print(f"  找到最新版本: {latest_version}. 正在为 '{original_line}' 添加版本锁定。") # New message
                    updated_lines.append(new_line)

            else: # latest_version is None (could not find latest) -> Keep original line
                print(f"  无法找到 {package_name} 的最新版本信息。保留原始行：'{original_line}'", file=sys.stderr) # Translated message
                updated_lines.append(original_line)

        # Write the updated content back to the file
        print(f"\n正在将更新后的内容写回 {requirements_file}...") # Translated message
        with open(requirements_file, 'w', encoding='utf-8') as f:
            for updated_line in updated_lines:
                f.write(updated_line + '\n') # Ensure each line ends with a newline

        print("Requirements 文件更新完成。") # Translated message

    except FileNotFoundError:
        print(f"错误：未找到文件 {requirements_file}。", file=sys.stderr) # Translated message
    except Exception as e:
        print(f"发生错误：{e}", file=sys.stderr) # Translated message

if __name__ == "__main__":
    # You can run this script directly.
    # Example: python update_requirements.py
    # Or pass a different requirements file name as an argument
    req_file = sys.argv[1] if len(sys.argv) > 1 else "requirements.txt"
    update_requirements(req_file)
